import os
import datetime
import json
import paramiko
from astropy.io import fits
from scp import SCPClient

class TemperatureCollector:
    def __init__(self, data_root):
        self.data_root = data_root
        self.camera_ids = self._load_camera_ids()
        self.username = 'gwac'
        self.password = 'gwac1234'
        
    def _load_camera_ids(self):
        """从camreaid文件加载相机ID列表"""
        with open('/data/temp_monitor/temp_web2/camreaid') as f:
            return [f'G{line.strip()}' for line in f if line.strip()]
    
    def _get_current_day_str(self):
        """获取当前UTC日期字符串"""
        current_utc = datetime.datetime.now(datetime.timezone.utc)
        return current_utc.strftime('%y%m%d')
    
    def _get_today_store_path_by_search(self, root_path):
        """本地实现getTodayStorePathBySearch函数"""
        today_store_path = ""
        start_obs = False

        template_file_name = "G005_054_241019"  
        template_file_name_len = len(template_file_name)  
        template_file_name_split_num = len(template_file_name.split("_"))  

        dir_pool = [root_path]  
        cur_date_str = self._get_current_day_str()

        for troot_dir in dir_pool:  
            stdin, stdout, stderr = self.ssh.exec_command(f'[ -d "{troot_dir}" ] && echo "exists" || echo "not exists"')
            if stdout.read().decode().strip() != "exists":
                continue
                
            stdin, stdout, stderr = self.ssh.exec_command(f'ls "{troot_dir}"')
            obs_days = stdout.read().decode().splitlines()
            
            for obs_day in obs_days:
                if (cur_date_str in obs_day and 
                    len(obs_day) == template_file_name_len and 
                    len(obs_day.split("_")) == template_file_name_split_num):
                    today_store_path = f"{troot_dir}/{obs_day}"
                    start_obs = True
                    break
            if start_obs:
                break

        return start_obs, today_store_path

    def _get_remote_path(self, camera_id, ssh):
        """获取相机服务器上的当天观测目录"""
        self.ssh = ssh  # 保存ssh连接供本地函数使用
        root_path = "/data"
        found, path = self._get_today_store_path_by_search(root_path)
        return path if found else None
    
    def _read_fits_temperature(self, filepath):
        """读取FITS文件温度信息"""
        try:
            with fits.open(filepath) as hdul:
                if filepath.endswith('.fit.fz'):
                    hdu = hdul[1]  # .fit.fz文件使用第1个头
                else:
                    hdu = hdul[0]  # .fit文件使用第0个头
                
                header = hdu.header
                return {
                    'TEMPSET': float(header.get('TEMPSET', 0)),
                    'TEMPACT': float(header.get('TEMPACT', 0)),
                    'DATE-OBS': header.get('DATE-OBS', ''),
                    'TIME-OBS': header.get('TIME-OBS', '')
                }
        except Exception as e:
            print(f"Error reading {filepath}: {e}")
            return None
    
    def _save_temperature_data(self, date, data):
        """保存温度数据到JSON文件"""
        # 创建日期子目录
        date_dir = os.path.join(self.data_root, date)
        os.makedirs(date_dir, exist_ok=True)
        
        # 设置输出文件路径
        output_file = os.path.join(date_dir, f"{date}_temp.json")
        
        # 合并现有数据(如果有)
        existing_data = {}
        if os.path.exists(output_file):
            with open(output_file) as f:
                existing_data = json.load(f)
        
        # 更新数据
        for cam_id, temp_data in data.items():
            if cam_id not in existing_data:
                existing_data[cam_id] = []
            existing_data[cam_id].append(temp_data)
        
        # 保存数据
        with open(output_file, 'w') as f:
            json.dump(existing_data, f, indent=2)
        
        # 更新日志中的文件路径显示
        return os.path.join(date, f"{date}_temp.json")
    
    def collect_temperatures(self):
        """从所有相机收集温度数据"""
        today = self._get_current_day_str()
        collected_data = {}
        total_cameras = len(self.camera_ids)
        success_count = 0
        
        print(f"\n=== 开始采集温度数据 ({today}) ===")
        print(f"共需采集 {total_cameras} 台相机数据")
        
        for idx, camera_id in enumerate(self.camera_ids, 1):
            camera_num = camera_id[1:]
            ip = f'172.28.2.{camera_num}'
            
            print(f"\n[{idx}/{total_cameras}] 正在采集 {camera_id} (IP: {ip})...")
            
            try:
                # 连接相机服务器
                print(" - 正在建立SSH连接...", end=' ')
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                ssh.connect(ip, username=self.username, password=self.password)
                print("成功")
                
                # 查找当天观测目录
                print(" - 正在查找当天观测目录...", end=' ')
                remote_path = self._get_remote_path(camera_id, ssh)
                if not remote_path:
                    print("失败: 未找到当天观测目录")
                    continue
                print(f"成功: {remote_path}")
                
                # 获取最新的FITS文件
                print(" - 正在查找最新FITS文件...", end=' ')
                stdin, stdout, stderr = ssh.exec_command(
                    f'ls -t {remote_path}/*.fit* | head -n 1'
                )
                latest_fits = stdout.read().decode().strip()
                
                if not latest_fits:
                    print("失败: 未找到FITS文件")
                    continue
                print(f"成功: {os.path.basename(latest_fits)}")
                
                # 检查本地是否已有该文件
                local_dir = os.path.join(self.data_root, today, camera_id)
                os.makedirs(local_dir, exist_ok=True)
                local_path = os.path.join(local_dir, os.path.basename(latest_fits))
                
                if not os.path.exists(local_path):
                    # 复制文件到本地
                    print(" - 正在复制文件到本地...", end=' ')
                    with SCPClient(ssh.get_transport()) as scp:
                        scp.get(latest_fits, local_path)
                    print("成功")
                    
                    # 读取温度数据
                    print(" - 正在读取温度数据...", end=' ')
                    temp_data = self._read_fits_temperature(local_path)
                    if temp_data:
                        collected_data[camera_id] = {
                            'TEMPSET': temp_data['TEMPSET'],
                            'TEMPACT': temp_data['TEMPACT'],
                            'time': f"{temp_data['DATE-OBS']} {temp_data['TIME-OBS']}"
                        }
                        success_count += 1
                        print(f"成功: TEMPSET={temp_data['TEMPSET']}°C, TEMPACT={temp_data['TEMPACT']}°C")
                    else:
                        print("失败: 无法读取温度数据")
                else:
                    print(" - 文件已存在，跳过采集")
                
            except Exception as e:
                print(f"失败: {str(e)}")
            finally:
                if 'ssh' in locals():
                    ssh.close()
        
        if collected_data:
            saved_path = self._save_temperature_data(today, collected_data)
            print(f"\n=== 采集完成 ===")
            print(f"成功采集 {success_count}/{total_cameras} 台相机数据")
            print(f"数据已保存到: {os.path.join(self.data_root, saved_path)}")
        else:
            print("\n=== 采集完成 ===")
            print("警告: 未采集到任何有效数据")